package io.gitee.skyarthur1987.leporidae.delayqueue.core;

import com.google.common.collect.Lists;
import io.gitee.skyarthur1987.leporidae.delayqueue.api.DelayQueue;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.BoundZSetOperations;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.util.Assert;

import java.util.Date;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * 基于 redis zset 结构实现的延迟队列
 * Created by skyarthur on 2023/1/26
 */
public class ZSetDelayQueue implements DelayQueue<String> {
    private final StringRedisTemplate stringRedisTemplate;

    public ZSetDelayQueue(StringRedisTemplate stringRedisTemplate) {
        this.stringRedisTemplate = stringRedisTemplate;
    }

    @Override
    public boolean add(String task, String topic, int delayTime, TimeUnit timeUnit) {
        Assert.notNull(task, "delay queue add: element cannot be null");
        Assert.notNull(topic, "delay queue add: topic cannot be null");
        Assert.isTrue(delayTime > 0, "delay queue add: delayTime need > 0");
        Assert.notNull(timeUnit, "delay queue add: timeUnit cannot be null");

        BoundZSetOperations<String, String> operations = stringRedisTemplate.boundZSetOps(topic);
        operations.add(task, System.currentTimeMillis() + timeUnit.toMillis(delayTime));
        return true;
    }

    @Override
    public boolean add(String task, String topic, Date executeTime) {
        Assert.notNull(task, "delay queue add: task cannot be null");
        Assert.notNull(topic, "delay queue add: topic cannot be null");
        Assert.notNull(executeTime, "delay queue add: executeTime cannot be null");
        long executeTimestamp = executeTime.getTime();
        BoundZSetOperations<String, String> operations = stringRedisTemplate.boundZSetOps(topic);
        operations.add(task, executeTimestamp);
        return true;
    }

    @Override
    public List<String> poll(String topic, Long expireTime) {
        Assert.notNull(topic, "delay queue poll: topic cannot be null");
        Assert.notNull(expireTime, "delay queue poll: expireTime cannot be null");

        DefaultRedisScript<List> redisScript = new DefaultRedisScript<>();
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("redis-delay-queue-poll.lua")));
        redisScript.setResultType(List.class);
        return stringRedisTemplate.execute(redisScript, Lists.newArrayList(topic),
                String.valueOf(expireTime));
    }
}
